// smrs.c
#include "smrs.h"
#include "smrs_config.h"
#include "xgpu_render_tbr.h"
#include "xil_printf.h"
#include "xstatus.h"
#include "xil_cache.h"
#include <string.h> // For memset
#include <math.h>   // For fmaxf, fminf
#include "linalg.h" // 引入自定义的线性代数库

// =========================================================================================
// PS端数据类型定义和转换函数 (模拟 HLS ap_fixed/ap_uint 的行为)
// 这里的转换函数是 SMRS 模块的私有实现细节，不暴露在 smrs.h 中
// =========================================================================================

// ==================== fp_coord_t (ap_fixed<28, 18, AP_RND_CONV, AP_SAT>) ====================
// HLS定义：typedef ap_fixed<28, 18, AP_RND_CONV, AP_SAT> fp_coord_t;
const int SMRS_FP_COORD_W = 28; // Total bits
const int SMRS_FP_COORD_I = 18; // Integer bits (including sign)
const int SMRS_FP_COORD_F = (SMRS_FP_COORD_W - SMRS_FP_COORD_I); // Fractional bits = 10
const float SMRS_FP_COORD_SCALE_FACTOR = 1024.0f; // 2^10 = 1024.0f

// HLS ap_fixed<W,I> 内部存储为 W 位的有符号整数。
// 对于 28 位，其范围是 [-2^27, 2^27 - 1]。
const s64 SMRS_FP_COORD_MAX_INTERNAL_VAL = (1LL << (SMRS_FP_COORD_W - 1)) - 1;   // max positive: 134217727
const s64 SMRS_FP_COORD_MIN_INTERNAL_VAL = -(1LL << (SMRS_FP_COORD_W - 1));     // min negative: -134217728

/**
 * @brief 将 PS 端浮点数 (float) 转换为 HLS IP核期望的 28 位定点数内部整数表示 (s32)。
 *        模拟 ap_fixed 的舍入 (AP_RND_CONV, round half to even) 和饱和 (AP_SAT) 行为。
 *        这个内部整数值会被写入 DDR/AXI-Lite。
 * @param val 输入浮点数
 * @return 转换为 28 位定点数的内部整数表示 (s32)。
 */
static inline s32 smrs_float_to_fp_coord_internal(float val) {
    float scaled_val = val * SMRS_FP_COORD_SCALE_FACTOR;
    long long internal_val_before_sat = (long long)roundf(scaled_val);

    s64 internal_val = internal_val_before_sat; // 使用 s64 进行饱和检查
    if (internal_val > SMRS_FP_COORD_MAX_INTERNAL_VAL) internal_val = SMRS_FP_COORD_MAX_INTERNAL_VAL;
    if (internal_val < SMRS_FP_COORD_MIN_INTERNAL_VAL) internal_val = SMRS_FP_COORD_MIN_INTERNAL_VAL;

    return (s32)internal_val; // 返回 28 位有符号整数
}


// ==================== fp_z_buffer_t (ap_uint<16>) ====================
// HLS定义：typedef ap_fixed<16, 16> fp_z_buffer_t; (通常是 ap_uint<16>)
/**
 * @brief 将 PS 端浮点数 (例如归一化深度) 转换为 16 位无符号 Z Buffer 值 (0-65535)。
 *        模拟 HLS ap_uint<16> 的截断行为。
 * @param z_float 输入浮点数
 * @return 转换后的 16 位无符号 Z Buffer 值 (u16)
 */
static inline u16 smrs_float_to_fp_z_buffer(float z_float) {
    // 钳制到 [0, 65535] 范围
    float clamped_val = fmaxf(0.0f, fminf(z_float, 65535.0f));
    u16 result = (u16)floorf(clamped_val); // 截断取整
    return result;
}


// ==================== fp_color_channel_t (ap_uint<8>) ====================
// HLS定义：typedef ap_uint<8> fp_color_channel_t;
/**
 * @brief 将 PS 端浮点数或整数转换为 8 位无符号颜色通道值 (0-255)。
 *        模拟 HLS ap_uint<8> 的截断行为。
 * @param color_val 输入数值
 * @return 转换后的 8 位无符号颜色通道值 (u8)
 */
static inline u8 smrs_to_fp_color_channel(float color_val) {
    // 钳制到 [0, 255] 范围
    float clamped_val = fmaxf(0.0f, fminf(color_val, 255.0f));
    u8 result = (u8)floorf(clamped_val); // 截断取整
    return result;
}


// =========================================================================================
// DDR 数据打包函数 (处理 64 位 smrs_ddr_word_t)
// =========================================================================================

/**
 * @brief 将 SmrsTriangle 打包成 HLS IP 期望的 64 位 DDR 字。
 * @param tri_ps 指向 SmrsTriangle 结构体的指针。
 * @param output_ddr_words 存储打包数据的 DDR 字数组。
 * @param triangle_idx 当前三角形在数组中的索引，用于计算输出DDR字数组的起始位置。
 */
static void smrs_pack_raster_params_to_ddr(const SmrsTriangle* tri_ps, smrs_ddr_word_t* output_ddr_words, u32 triangle_idx) {
    // 每个三角形占据 SMRS_NUM_AXI_WORDS_PER_PACKET 个 64 位 DDR 字
    u32 current_packet_word_offset = triangle_idx * SMRS_NUM_AXI_WORDS_PER_PACKET;
    u32 word_idx = 0; // 当前数据包内的字索引 (0 to SMRS_NUM_AXI_WORDS_PER_PACKET-1)

    s32 internal_val_coord; // 28位定点坐标的内部整数表示
    u16 internal_val_z;     // 16位ZBuffer值
    smrs_ddr_word_t ddr_word_temp;
    
    // --- 顶点坐标 (6 words) ---
    internal_val_coord = smrs_float_to_fp_coord_internal(tri_ps->v0x);
    ddr_word_temp = (smrs_ddr_word_t)((s64)internal_val_coord);
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 0
    
    internal_val_coord = smrs_float_to_fp_coord_internal(tri_ps->v0y);
    ddr_word_temp = (smrs_ddr_word_t)((s64)internal_val_coord);
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 1
    
    internal_val_coord = smrs_float_to_fp_coord_internal(tri_ps->v1x);
    ddr_word_temp = (smrs_ddr_word_t)((s64)internal_val_coord);
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 2
    
    internal_val_coord = smrs_float_to_fp_coord_internal(tri_ps->v1y);
    ddr_word_temp = (smrs_ddr_word_t)((s64)internal_val_coord);
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 3
    
    internal_val_coord = smrs_float_to_fp_coord_internal(tri_ps->v2x);
    ddr_word_temp = (smrs_ddr_word_t)((s64)internal_val_coord);
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 4
    
    internal_val_coord = smrs_float_to_fp_coord_internal(tri_ps->v2y);
    ddr_word_temp = (smrs_ddr_word_t)((s64)internal_val_coord);
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 5
    
    // --- Z 值 (3 words) ---
    // ZBuffer 的深度测试逻辑通常为 `interpolated_z <= line_z_buffer[idx]` (Z值越小越近)。
    // SMRS_NUM_AXI_WORDS_PER_PACKET 
    // 所以 Z 值映射：原始 vNz_norm 是 [0, 1] 范围，0为近，1为远。
    // 转换为 [0, 65535] 后，65535为近，0为远 (Z值越大越近)。
    // ZBuffer 初始值清空为 0xFFFF (最远)，所以这里将 0-1 范围的归一化 Z 反向映射为 (1.0f - vNz_norm) * 65535.0f。
    internal_val_z = smrs_float_to_fp_z_buffer((1.0f - tri_ps->v0z_norm) * 65535.0f);
    ddr_word_temp = (smrs_ddr_word_t)internal_val_z;
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 6
    
    internal_val_z = smrs_float_to_fp_z_buffer((1.0f - tri_ps->v1z_norm) * 65535.0f);
    ddr_word_temp = (smrs_ddr_word_t)internal_val_z;
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 7
    
    internal_val_z = smrs_float_to_fp_z_buffer((1.0f - tri_ps->v2z_norm) * 65535.0f);
    ddr_word_temp = (smrs_ddr_word_t)internal_val_z;
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 8
    
    // --- 边缘方程系数 (9 words) ---
    internal_val_coord = smrs_float_to_fp_coord_internal(tri_ps->edge0_A);
    ddr_word_temp = (smrs_ddr_word_t)((s64)internal_val_coord);
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 9
    
    internal_val_coord = smrs_float_to_fp_coord_internal(tri_ps->edge0_B);
    ddr_word_temp = (smrs_ddr_word_t)((s64)internal_val_coord);
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 10
    
    internal_val_coord = smrs_float_to_fp_coord_internal(tri_ps->edge0_C);
    ddr_word_temp = (smrs_ddr_word_t)((s64)internal_val_coord);
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 11
    
    internal_val_coord = smrs_float_to_fp_coord_internal(tri_ps->edge1_A);
    ddr_word_temp = (smrs_ddr_word_t)((s64)internal_val_coord);
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 12
    
    internal_val_coord = smrs_float_to_fp_coord_internal(tri_ps->edge1_B);
    ddr_word_temp = (smrs_ddr_word_t)((s64)internal_val_coord);
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 13
    
    internal_val_coord = smrs_float_to_fp_coord_internal(tri_ps->edge1_C);
    ddr_word_temp = (smrs_ddr_word_t)((s64)internal_val_coord);
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 14
    
    internal_val_coord = smrs_float_to_fp_coord_internal(tri_ps->edge2_A);
    ddr_word_temp = (smrs_ddr_word_t)((s64)internal_val_coord);
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 15
    
    internal_val_coord = smrs_float_to_fp_coord_internal(tri_ps->edge2_B);
    ddr_word_temp = (smrs_ddr_word_t)((s64)internal_val_coord);
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 16
    
    internal_val_coord = smrs_float_to_fp_coord_internal(tri_ps->edge2_C);
    ddr_word_temp = (smrs_ddr_word_t)((s64)internal_val_coord);
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 17
    
    // --- 深度梯度 (2 words) ---
    // 注意：这里的 dzdx, dzdy 已经是在 `Smrs_CalculateDepthGradients` 中基于 [0, 65535] 范围的 Z 计算得到
    internal_val_coord = smrs_float_to_fp_coord_internal(tri_ps->dzdx);
    ddr_word_temp = (smrs_ddr_word_t)((s64)internal_val_coord);
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 18
    
    internal_val_coord = smrs_float_to_fp_coord_internal(tri_ps->dzdy);
    ddr_word_temp = (smrs_ddr_word_t)((s64)internal_val_coord);
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 19
    
    // --- 颜色 (3 words) ---
    // HLS IP 期望 R, G, B 各占据一个 64 位 DDR 字的低 8 位
    u8 final_r = smrs_to_fp_color_channel(tri_ps->r_f);
    u8 final_g = smrs_to_fp_color_channel(tri_ps->g_f);
    u8 final_b = smrs_to_fp_color_channel(tri_ps->b_f);
    
    ddr_word_temp = (smrs_ddr_word_t)final_r;
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 20 (Red)
    
    ddr_word_temp = (smrs_ddr_word_t)final_g;
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 21 (Green)
    
    ddr_word_temp = (smrs_ddr_word_t)final_b;
    output_ddr_words[current_packet_word_offset + word_idx++] = ddr_word_temp; // word 22 (Blue)
    
    // 检查打包的字数是否匹配
    if (word_idx != SMRS_NUM_AXI_WORDS_PER_PACKET) {
        xil_printf("SMRS ERROR: Mismatch in expected words per packet! Expected %u, but wrote %u for triangle %u\r\n", SMRS_NUM_AXI_WORDS_PER_PACKET, word_idx, triangle_idx);
    }
}


// =========================================================================================
// SMRS 模块公共 API 实现
// =========================================================================================

int Smrs_Init(SmrsContext* pContext, const VdmaDisplay* vdma_display_ptr) {
    if (pContext == NULL || vdma_display_ptr == NULL) {
        xil_printf("SMRS ERROR: Smrs_Init received NULL context or VdmaDisplay pointer.\r\n");
        return SMRS_FAILURE;
    }

    pContext->render_width = VdmaDisplay_getHlsRenderWidth(vdma_display_ptr);
    pContext->render_height = VdmaDisplay_getHlsRenderHeight(vdma_display_ptr);
    pContext->hls_single_frame_size = VdmaDisplay_getHlsSingleFrameSize(vdma_display_ptr);

    // 初始化 GPU_Render_TBR IP
    XGpu_render_tbr_Config *GpuRenderTbrCfgPtr;
    
    GpuRenderTbrCfgPtr = XGpu_render_tbr_LookupConfig(SMRS_GPU_RENDER_TBR_BASE_ADDR);
    if (GpuRenderTbrCfgPtr == NULL) {
        xil_printf("SMRS ERROR: XGpu_render_tbr_LookupConfig failed! BaseAddress=0x%x\r\n", (unsigned int)SMRS_GPU_RENDER_TBR_BASE_ADDR);
        return SMRS_FAILURE;
    }
    
    int status = XGpu_render_tbr_CfgInitialize(&pContext->GpuRenderTbrInstance, GpuRenderTbrCfgPtr);
    if (status != XST_SUCCESS) {
        xil_printf("SMRS ERROR: XGpu_render_tbr_CfgInitialize failed! Status=%d.\r\n", status);
        return SMRS_FAILURE;
    }
    xil_printf("SMRS: GPU_Render_TBR IP initialized successfully.\r\n");
    
    // 自检：检查 GPU IP 是否处于就绪状态
    if (!XGpu_render_tbr_IsReady(&pContext->GpuRenderTbrInstance)) {
        xil_printf("SMRS WARNING: GPU_Render_TBR IP is not ready after initialization.\r\n");
    } else {
        xil_printf("SMRS: GPU_Render_TBR IP is ready.\r\n");
    }

    // 分配 PS 端用于打包三角形数据的 DDR 内存
    pContext->params_buffer = (smrs_ddr_word_t*)SMRS_GPU_PARAMS_BUFFER_BASE_ADDR;
    if (pContext->params_buffer == NULL) {
        xil_printf("SMRS ERROR: Failed to allocate params_buffer at 0x%x.\r\n", (unsigned int)SMRS_GPU_PARAMS_BUFFER_BASE_ADDR);
        return SMRS_FAILURE;
    }
    xil_printf("SMRS: Allocated params_buffer at 0x%x, size %u bytes.\r\n", pContext->params_buffer, (unsigned int)SMRS_GPU_PARAMS_TOTAL_BYTE_SIZE);
    
    // Z-buffer 分配
    pContext->zbuffer_raw_ptr = (u16*)SMRS_GPU_ZBUFFER_BASE_ADDR;
    pContext->zbuffer_size_bytes = pContext->render_width * pContext->render_height * sizeof(u16);
    
    if (pContext->zbuffer_raw_ptr == NULL) {
        xil_printf("SMRS ERROR: Failed to allocate zbuffer at 0x%x.\r\n", (unsigned int)SMRS_GPU_ZBUFFER_BASE_ADDR);
        return SMRS_FAILURE;
    }
    xil_printf("SMRS: Allocated zbuffer at 0x%x, size %u bytes.\r\n", pContext->zbuffer_raw_ptr, (unsigned int)pContext->zbuffer_size_bytes);

    return SMRS_SUCCESS;
}


int Smrs_ClearFrameAndZBuffer(SmrsContext* pContext, UINTPTR current_hls_framebuffer_addr) {
    if (pContext == NULL || current_hls_framebuffer_addr == 0) {
        xil_printf("SMRS ERROR: Smrs_ClearFrameAndZBuffer received invalid parameters.\r\n");
        return SMRS_FAILURE;
    }

    // 清空目标 HLS 帧缓冲区为默认背景色 (深灰色 BGRX: 0x00202020)
    u32 default_color_slot = Smrs_PackRgbTo32BitSlot(0x20, 0x20, 0x20); // R=32, G=32, B=32
    u32* hls_framebuffer_ptr_u32 = (u32*)current_hls_framebuffer_addr;
    for (u32 i = 0; i < (pContext->render_width * pContext->render_height); ++i) {
        hls_framebuffer_ptr_u32[i] = default_color_slot;
    }
    Xil_DCacheFlushRange((INTPTR)current_hls_framebuffer_addr, pContext->hls_single_frame_size);

    // 重置 ZBuffer 为最远值 (0xFFFF for each 16-bit Z value)
    Xil_DCacheInvalidateRange((INTPTR)pContext->zbuffer_raw_ptr, pContext->zbuffer_size_bytes); // 确保 CPU 读取 ZBuffer 为最新状态
    for (u32 i = 0; i < (pContext->render_width * pContext->render_height); ++i) {
        pContext->zbuffer_raw_ptr[i] = 0xFFFF; // 初始Z值设置为最大 (最远)
    }
    Xil_DCacheFlushRange((INTPTR)pContext->zbuffer_raw_ptr, pContext->zbuffer_size_bytes); // 清零后 Flush 回 DDR

    return SMRS_SUCCESS;
}


int Smrs_RenderTriangles(SmrsContext* pContext, const SmrsTriangle* triangles, u32 num_triangles, UINTPTR target_hls_framebuffer_addr) {
    if (pContext == NULL || triangles == NULL || num_triangles == 0 || target_hls_framebuffer_addr == 0) {
        // 允许 num_triangles 为 0，表示没有三角形需要渲染，直接返回成功
        if (num_triangles == 0) return SMRS_SUCCESS;

        xil_printf("SMRS ERROR: Smrs_RenderTriangles received invalid parameters.\r\n");
        return SMRS_FAILURE;
    }

    // 将所有有效三角形的参数打包并写入 params_buffer
    // params_buffer 存储了当前帧所有经过处理的三角形数据
    for (u32 i = 0; i < num_triangles; ++i) {
        smrs_pack_raster_params_to_ddr(&triangles[i], pContext->params_buffer, i);
    }
    Xil_DCacheFlushRange((INTPTR)pContext->params_buffer, num_triangles * SMRS_NUM_AXI_WORDS_PER_PACKET * sizeof(smrs_ddr_word_t));
    
    // 分批次配置并启动 GPU_Render_TBR IP
    u32 current_triangle_offset_batch = 0;
    while (current_triangle_offset_batch < num_triangles) {
        u32 num_triangles_in_current_batch = fminf((float)SMRS_MAX_NUM_TRIANGLES, (float)(num_triangles - current_triangle_offset_batch));
        
        // 计算当前批次三角形参数在 params_buffer 中的实际物理地址
        UINTPTR batch_params_base_addr = SMRS_GPU_PARAMS_BUFFER_BASE_ADDR + (current_triangle_offset_batch * SMRS_NUM_AXI_WORDS_PER_PACKET * sizeof(smrs_ddr_word_t));

        XGpu_render_tbr_Set_params_buffer_base_addr(&pContext->GpuRenderTbrInstance, (u64)batch_params_base_addr);
        XGpu_render_tbr_Set_num_triangles_to_process(&pContext->GpuRenderTbrInstance, num_triangles_in_current_batch);
        XGpu_render_tbr_Set_max_x_res(&pContext->GpuRenderTbrInstance, pContext->render_width);
        XGpu_render_tbr_Set_max_y_res(&pContext->GpuRenderTbrInstance, pContext->render_height);
        XGpu_render_tbr_Set_zbuffer_ddr_raw_ptr(&pContext->GpuRenderTbrInstance, (u64)SMRS_GPU_ZBUFFER_BASE_ADDR);
        XGpu_render_tbr_Set_framebuffer_ddr_raw_ptr(&pContext->GpuRenderTbrInstance, (u64)target_hls_framebuffer_addr);
        XGpu_render_tbr_Start(&pContext->GpuRenderTbrInstance);

        // 等待当前批次 GPU_Render_TBR IP 完成 (使用 Xil_GPU_Render_TBR_IsDone 而不是 FreeRTOS 的 portYIELD)
        while (!XGpu_render_tbr_IsDone(&pContext->GpuRenderTbrInstance)) {
            // 在 FreeRTOS 环境中，这里可以调用 portYIELD() 允许其他任务运行
            // 如果不在 FreeRTOS 环境，则会忙等
            #ifdef FREERTOS_USED
            portYIELD();
            #endif
        }
        current_triangle_offset_batch += num_triangles_in_current_batch;
    }
    
    return SMRS_SUCCESS;
}


u32 Smrs_PackRgbTo32BitSlot(u8 r, u8 g, u8 b) {
    u32 packed_color = 0;
    packed_color |= (u32)b << 0;  // B (0-7)
    packed_color |= (u32)g << 8;  // G (8-15)
    packed_color |= (u32)r << 16; // R (16-23)
    return packed_color; // 高 8 位是 0
}


void Smrs_CalculateEdgeEquations(
    float v0x, float v0y, float v1x, float v1y, float v2x, float v2y,
    float* A0, float* B0, float* C0,
    float* A1, float* B1, float* C1,
    float* A2, float* B2, float* C2
) {
    *A0 = v0y - v1y;
    *B0 = v1x - v0x;
    *C0 = v0x * v1y - v1x * v0y;

    *A1 = v1y - v2y;
    *B1 = v2x - v1x;
    *C1 = v1x * v2y - v2x * v1y;
    
    *A2 = v2y - v0y;
    *B2 = v0x - v2x;
    *C2 = v2x * v0y - v0x * v2y;
}


void Smrs_CalculateDepthGradients(
    float v0x, float v0y, float v0z, // v0z is screen-space Z (0-65535)
    float v1x, float v1y, float v1z,
    float v2x, float v2y, float v2z,
    float* dzdx, float* dzdy
) {
    float dx1 = v1x - v0x;
    float dy1 = v1y - v0y;
    float dz1 = v1z - v0z;

    float dx2 = v2x - v0x;
    float dy2 = v2y - v0y;
    float dz2 = v2z - v0z;
    
    float det = dx1 * dy2 - dx2 * dy1;
    
    if (fabsf(det) < 1e-6f) { // 防止除以零
        *dzdx = 0.0f;
        *dzdy = 0.0f;
        return;
    }
    
    float inv_det = 1.0f / det;
    *dzdx = (dz1 * dy2 - dz2 * dy1) * inv_det;
    *dzdy = (dx1 * dz2 - dx2 * dy1) * inv_det;
}

